#ifndef SATLAB_H
#define SATLAB_H

#include <vector>
#include <cstdlib>
#include <cstring>
#include <cstdio>
#include <unordered_map>
#include <string>
#include <iostream>
#include <fstream>
#include <filesystem>
#include <algorithm>
#include <lzma.h>

#include <vector>
#include <opencv2/opencv.hpp>
#include <string>

#include "dexvec.h"
#include "dexmat.h"

constexpr double MIN_OPACITY = 0.8;

struct Pixel {
    uint64_t pos = 0;
    uint64_t neg = 0;
};

struct Block {
    std::vector<std::vector<Pixel>> pixels;  // 2D grid of pixels within the block
    size_t varStart, varEnd;  // Variable range for this block
    size_t clauseStart, clauseEnd;  // Clause range for this block
};

class SATLAB {
public:
    struct Seg {
        size_t ptr;
        size_t dex;
    };

    DexMat mat;
    std::vector<std::vector<Pixel>> scope;
    std::vector<Seg> claSeg;
    std::vector<Block> currentBlockRow;  // New: store only current row of blocks
    size_t varNum = 0;
    size_t claNum = 0;
    cv::Vec3b posColor;
    cv::Vec3b negColor;
    size_t image_w;
    size_t image_h;
    size_t blockVarSize;  // New: variables per block
    size_t blockClauseSize;  // New: clauses per block
    size_t currentBlockRowIndex;  // New: current row being processed

    void calcBlockSize();

    SATLAB(int image_w = 512, int image_h = 512);

    void allocate(int varNum, int claNum);
    // void sortClause();
    void updateAllStats();
    void updateAllStatsForBlockRow(size_t targetBlockRow);
    cv::Mat& createImage(cv::Mat& image, double colorThreshRatio = 0.5);
    cv::Mat& createBinScopeImage(cv::Mat& image, size_t binIndex, double colorThreshRatio = 0.5);
    void createAllBinScopeImages(const std::string& outputDir, double colorThreshRatio = 0.5);
    cv::Mat& createBlockImage(cv::Mat& image, size_t blockIndex, double colorThreshRatio = 0.5);
    void createCurrentBlockRowImages(const std::string& outputDir, size_t rowIndex, double colorThreshRatio = 0.5);
    void initializeBlockRow(size_t rowIndex);
    void clearBlockRow();
    void processBlockRowIfComplete(const std::string& outputDir, size_t currentClauseIndex);
    cv::Mat& createHoughTransformImage(cv::Mat& image, const cv::Mat& scopeImage);
    cv::Mat& createFFTImage(cv::Mat& image, const cv::Mat& scopeImage);
    cv::Mat& createMorphImage(cv::Mat& image, const cv::Mat& scopeImage);

    void showScope();

    void parseFromMemory(char* data);
    void parseCompressed(const char* filename);
    void parsePlainText(const char* filename);
    void parse(const char* filename);

    // Serialization for scope array
    void saveScopeArray(const std::string& filename, const std::vector<std::vector<Pixel>>& scope);
    void loadScopeArray(const std::string& filename, std::vector<std::vector<Pixel>>& scope);
    // Serialization for DexMat
    void saveDexMat(const std::string& filename, const DexMat& mat);
    void loadDexMat(const std::string& filename, DexMat& mat);

    // Generate direct and pooled images from a file (handles compressed and uncompressed)
    void createByteImages(const std::string& filename, cv::Mat& directImg, cv::Mat& minImg, cv::Mat& maxImg, cv::Mat& meanImg);

private:
    // ... existing code ...
};

#endif 